package com.smh.aiagent.controller;

import com.smh.aiagent.agent.SmhAgent;
import com.smh.aiagent.app.LoveApp;
import com.smh.aiagent.model.vo.ResultVO;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;

/**
 * @author shiminghui
 * @date 2025/6/9 20:05
 * @description: TODO
 */
@RestController
@RequestMapping("/ai")
public class AIController {

    @Resource
    private LoveApp loveApp;

    @Resource
    private ChatModel dashscopeChatModel;

    @Resource
    private ToolCallback[] allTools;


    @GetMapping("/love_app/chat/sync")
    public ResultVO<String> doChatWithLoveAppSync(@RequestParam("message") String userPrompt, @RequestParam("conversationId") String conversationId) {
        String s = loveApp.doChatWithSync(userPrompt, conversationId);
        return ResultVO.success(s);
    }

    @GetMapping(value = "love_app/chat/sse1", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> doChatWithLoveAppSse(@RequestParam("message") String userPrompt, @RequestParam("conversationId") String conversationId) {
        return loveApp.doChatByStream(userPrompt, conversationId);
    }

    @GetMapping(value = "/love_app/chat/sse2")
    public Flux<ServerSentEvent<String>> doChatWithLoveAppSSE(String message, String chatId) {
        return loveApp.doChatByStream(message, chatId)
                .map(chunk -> ServerSentEvent.<String>builder()
                        .data(chunk)
                        .build());
    }

    @GetMapping("/love_app/chat/sse/emitter")
    public SseEmitter doChatWithLoveAppSseEmitter(String message, String chatId) {
        // 创建一个超时时间较长的 SseEmitter
        SseEmitter emitter = new SseEmitter(180000L); // 3分钟超时
        // 获取 Flux 数据流并直接订阅
        loveApp.doChatByStream(message, chatId)
                .subscribe(
                        // 处理每条消息
                        chunk -> {
                            try {
                                emitter.send(chunk);
                            } catch (IOException e) {
                                emitter.completeWithError(e);
                            }
                        },
                        // 处理错误
                        emitter::completeWithError,
                        // 处理完成
                        emitter::complete
                );
        // 返回emitter
        return emitter;
    }

    @GetMapping("/agent/chat/sse/emitter")
    public SseEmitter doChatWithAgent(String message) {
        return SmhAgent.builder()
                .chatClient(ChatClient.builder(dashscopeChatModel).build())
                .availableTools(allTools)
                .build().runStream(message);
    }


}
